import mysql.connector
import ksycopg2
from ksycopg2.extras import RealDictCursor
from mysql.connector import Error
import os
client_table_name = os.getenv('DB_CLIENT_DATATABLE', 'client_data')
data_source = None
DATA_DEFAULT_SOURCE = 'mysql'
def create_connection():
    try:
        global data_source
        data_source = os.getenv('DB_DATA_SOURCE', DATA_DEFAULT_SOURCE)
        if data_source == DATA_DEFAULT_SOURCE:
            connection = mysql.connector.connect(
                host=os.getenv('DB_HOST', '10.21.231.40'),  # 使用环境变量，默认值为 '10.21.231.40'
                port=int(os.getenv('DB_PORT', 3306)),  # 使用环境变量，默认值为 3306
                user=os.getenv('DB_USER', 'client_user'),  # 使用环境变量，默认值为 'agg_user'
                password=os.getenv('DB_PASSWORD', 'client_user_password'),  # 使用环境变量，默认值为 'agg_user_password'
                database=os.getenv('DB_DATABASE', 'client_database')  # 使用环境变量，默认值为 'agg_database'
            )
            if connection.is_connected():
                print('Connected to MySQL database')
                return connection

        # 连接人大金仓数据库
        else:
            connection = ksycopg2.connect(
                host=os.getenv('DB_HOST', '10.161.43.163'),  # 使用环境变量，默认值为 '10.21.231.40'
                port=os.getenv('DB_PORT', '54321'),  # 使用环境变量，默认值为 '54321'
                user=os.getenv('DB_USER', 'system'),  # 使用环境变量，默认值为 'trace_user'
                password=os.getenv('DB_PASSWORD', '123456'),  # 使用环境变量，默认值为 'trace_user_password'
                database=os.getenv('DB_DATABASE', 'fed_database')  # 使用环境变量，默认值为 'trace_database'
            )
            if connection:
                print('Connected to KingbaseES database')
                return connection
    except Error as e:
        print(f"Error: {e}")
        return None

def create_client_data_table(connection):
    cursor = connection.cursor()
    if data_source == DATA_DEFAULT_SOURCE:
        cursor.execute(f"""
            CREATE TABLE IF NOT EXISTS {client_table_name} (
                id INT AUTO_INCREMENT PRIMARY KEY,
                client_id VARCHAR(255) NOT NULL,
                local_model LONGTEXT,
                local_BFV LONGBLOB,
                local_sign1 TEXT,
                local_sign2 TEXT,
                sk_sign_1 TEXT,
                pk_sign_1 TEXT,
                sk_sign_2 TEXT,
                pk_sign_2 TEXT,
                client_num INT,
                userid VARCHAR(255),
                username VARCHAR(255),
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
    else :
        cursor.execute(f"""
                    CREATE TABLE IF NOT EXISTS {client_table_name} (
                        id SERIAL PRIMARY KEY,
                        client_id VARCHAR(255) NOT NULL,
                        local_model LONGTEXT,
                        local_BFV BLOB,
                        local_sign1 TEXT,
                        local_sign2 TEXT,
                        sk_sign_1 TEXT,
                        pk_sign_1 TEXT,
                        sk_sign_2 TEXT,
                        pk_sign_2 TEXT,
                        client_num INT,
                        userid VARCHAR(255),
                        username VARCHAR(255),
                        timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                """)
    connection.commit()



def insert_or_update_client_data(connection, client_id, local_model=None, local_BFV=None, local_sign1=None,
                                 local_sign2=None, sk_sign_1=None, pk_sign_1=None, sk_sign_2=None, pk_sign_2=None,
                                 client_num=None, userid=None, username=None):
    print(client_table_name)
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT local_model, local_BFV, local_sign1, local_sign2, sk_sign_1, pk_sign_1, sk_sign_2, pk_sign_2, client_num, userid, username
        FROM {client_table_name} WHERE client_id = %s
    """, (client_id,))
    record = cursor.fetchone()

    if record:
        local_model = local_model if local_model is not None else record[0]
        local_BFV = local_BFV if local_BFV is not None else record[1]
        local_sign1 = local_sign1 if local_sign1 is not None else record[2]
        local_sign2 = local_sign2 if local_sign2 is not None else record[3]
        sk_sign_1 = sk_sign_1 if sk_sign_1 is not None else record[4]
        pk_sign_1 = pk_sign_1 if pk_sign_1 is not None else record[5]
        sk_sign_2 = sk_sign_2 if sk_sign_2 is not None else record[6]
        pk_sign_2 = pk_sign_2 if pk_sign_2 is not None else record[7]
        client_num = client_num if client_num is not None else record[8]
        userid = userid if userid is not None else record[9]
        username = username if username is not None else record[10]

        cursor.execute(f"""
            UPDATE {client_table_name}
            SET local_model = %s, local_BFV = %s, local_sign1 = %s, local_sign2 = %s, 
                sk_sign_1 = %s, pk_sign_1 = %s, sk_sign_2 = %s, pk_sign_2 = %s, client_num = %s,
                userid = %s, username = %s, timestamp = CURRENT_TIMESTAMP
            WHERE client_id = %s
        """, (local_model, local_BFV, local_sign1, local_sign2, sk_sign_1, pk_sign_1, sk_sign_2, pk_sign_2, client_num, userid, username, client_id))
    else:
        cursor.execute(f"""
            INSERT INTO {client_table_name} (client_id, local_model, local_BFV, local_sign1, local_sign2, sk_sign_1, pk_sign_1, sk_sign_2, pk_sign_2, client_num, userid, username)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """, (client_id, local_model, local_BFV, local_sign1, local_sign2, sk_sign_1, pk_sign_1, sk_sign_2, pk_sign_2, client_num, userid, username))
    connection.commit()


# def get_client_id(connection):
#     cursor = connection.cursor()
#     cursor.execute(f"""
#         SELECT client_id FROM {client_table_name}
#     """)
#     return cursor.fetchone()
def get_client_id(connection, client_id):
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT client_id FROM {client_table_name} WHERE client_id = %s
    """, (client_id,))
    result = cursor.fetchone()
    print(result)
    # 如果查询有结果，返回 True；否则返回 False
    return result is not None


def get_client_data(connection, client_id):
    if data_source == DATA_DEFAULT_SOURCE:
        cursor = connection.cursor(dictionary=True)
        try:
            cursor.execute(f"""
                SELECT local_BFV, local_sign1, local_sign2, pk_sign_1, pk_sign_2, client_num, userid, username
                FROM {client_table_name} WHERE client_id = %s
            """, (client_id,))
            result = cursor.fetchone()
            print("local_sign1", result['local_sign1'])
            print("local_sign2", result['local_sign2'])
            print("pk_sign_1", result['pk_sign_1'])
            print("pk_sign_2", result['pk_sign_2'])
            print("local_BFV", type(result['local_BFV']))
            print("local_BFV_bytes", result['local_BFV'][:20])
            print("local_BFV_bytes", result['local_BFV'].hex()[:40])

            if result:
                result['client_id'] = client_id
            return result if result else {}
        except Error as e:
            print(f"Error: {e}")
            return {}
    else:
        cursor = connection.cursor(cursor_factory=RealDictCursor)
        try:
            cursor.execute(f"""
                SELECT local_BFV, local_sign1, local_sign2, pk_sign_1, pk_sign_2, client_num, userid, username
                FROM {client_table_name} WHERE client_id = %s
            """, (client_id,))
            result = cursor.fetchone()
            print("local_sign1",result['local_sign1'])
            print("local_sign2", result['local_sign2'])
            print("pk_sign_1", result['pk_sign_1'])
            print("pk_sign_2", result['pk_sign_2'])
            print("local_BFV_bytes", result['local_BFV'][:20])
            print("local_BFV_bytes", result['local_BFV'])
            print("local_BFV_bytes", result['local_BFV'].hex()[:40])

            if result:
                result['client_id'] = client_id
            return result if result else {}
        except Error as e:
            print(f"Error: {e}")
            return {}


def get_local_sign1_data(connection, client_id):
    cursor = connection.cursor(dictionary=True)
    try:
        cursor.execute(f"""
            SELECT local_sign1 FROM {client_table_name} WHERE client_id = %s
        """, (client_id,))
        result = cursor.fetchone()
        # result['client_id'] = client_id
        return result if result else {}
    except Error as e:
        print(f"Error: {e}")
        return {}

# def drop_client_data_table():
#     try:
#         connection = create_connection()
#         cursor = connection.cursor()
#         cursor.execute(f"""DROP TABLE IF EXISTS {client_table_name}""")
#         connection.commit()
#     except Error as e:
#         print(f"Database error: {e}")
#     finally:
#         if connection.is_connected():
#             cursor.close()
#             connection.close()

def delete_client_data(client_id):
    try:
        connection = create_connection()
        cursor = connection.cursor()
        cursor.execute(f"""DELETE FROM {client_table_name} WHERE client_id = %s""", (client_id,))
        connection.commit()
        print(cursor.rowcount,type(cursor.rowcount))
        return cursor.rowcount  # 返回删除的行数
    except Error as e:
        print(f"Database error: {e}")
        return None  # 返回 None 以表示出现错误
    finally:
        if connection:
            cursor.close()
            connection.close()

